/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * openUBMC is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#ifndef MC_BUILTINS_LAMBDA_H
#define MC_BUILTINS_LAMBDA_H

#include <array>
#include <mc/exception.h>
#include <mc/variant.h>
#include <mc/variant/variant_extension.h>
#include <mcpy/builtins/cast.h>
#include <tuple>
#include <utility>

namespace mcpy {

namespace detail {

MC_API mc::variants fill_defaults(
    const mc::variants& args,
    const mc::variants& defaults,
    size_t              first_default_index,
    size_t              total_param_count);

template <typename Sig>
struct invoke_helper;

template <typename Ret, typename... Args>
struct invoke_helper<Ret(Args...)> {
    static constexpr size_t arg_count = sizeof...(Args);

    template <typename Func>
    static mc::variant invoke(Func&& func, const mc::variants& args, const std::array<const char*, arg_count>& param_names) {
        return invoke_impl<Func>(
            std::forward<Func>(func),
            args,
            param_names,
            std::make_index_sequence<arg_count>());
    }

    // 兼容旧接口（无参数名）
    template <typename Func>
    static mc::variant invoke(Func&& func, const mc::variants& args) {
        std::array<const char*, arg_count> default_names;
        for (size_t i = 0; i < arg_count; ++i) {
            default_names[i] = "lambda";
        }
        return invoke(std::forward<Func>(func), args, default_names);
    }

private:
    template <typename Func, size_t... I>
    static mc::variant invoke_impl(Func&& func, const mc::variants& args, const std::array<const char*, arg_count>& param_names, std::index_sequence<I...>) {
        if constexpr (std::is_void_v<Ret>) {
            func(mc::detail::convert_arg<mc::traits::remove_cvref_t<Args>>(
                param_names[I],
                args[I])...);
            return {};
        } else {
            return func(mc::detail::convert_arg<mc::traits::remove_cvref_t<Args>>(
                param_names[I], args[I])...);
        }
    }
};

} // namespace detail

// ============================================================================
// Lambda Extension Base
// ============================================================================

/**
 * @brief Lambda 扩展基类
 */
class MC_API lambda_extension_base : public mc::variant_extension_base {
public:
    explicit lambda_extension_base() {
    }
    virtual ~lambda_extension_base() = default;

    mc::variant call(const mc::variants& args);

    std::string_view        get_type_name() const override;
    mc::extension_type_info get_type_info() const override;
    bool                    equals(const mc::variant_extension_base& other) const override;
    std::size_t             hash() const override;

    mc::shared_ptr<mc::variant_extension_base> copy() const override;
    mc::shared_ptr<mc::variant_extension_base> deep_copy(
        mc::detail::copy_context* ctx = nullptr) const override;

protected:
    virtual mc::variant call_impl(const mc::variants& args) = 0;
    virtual size_t      get_arg_count() const               = 0;

    // 获取最小参数数量（派生类应重写以支持默认参数）
    virtual size_t get_min_arg_count() const {
        return get_arg_count(); // 默认：所有参数都是必需的
    }
};

// ============================================================================
// Lambda 基类
// ============================================================================

template <typename Derived, typename Sig>
class lambda_base;

template <typename Derived, typename Ret, typename... Args>
class lambda_base<Derived, Ret(Args...)> : public lambda_extension_base {
public:
    using return_type = Ret;
    using arg_types   = std::tuple<Args...>;

    static constexpr size_t arg_count = sizeof...(Args);

    lambda_base() = default;

    // 设置默认值数组
    void set_defaults(mc::variants&& defaults, size_t first_default_index) {
        defaults_            = std::move(defaults);
        first_default_index_ = first_default_index;
    }

    mc::variant call_impl(const mc::variants& args) override {
        // 如果有默认值且参数不足，填充默认值
        if (!defaults_.empty() && args.size() < sizeof...(Args)) {
            mc::variants full_args = detail::fill_defaults(
                args, defaults_, first_default_index_, sizeof...(Args));

            // 使用填充后的参数调用
            return detail::invoke_helper<Ret(Args...)>::invoke(
                *static_cast<Derived*>(this), full_args, Derived::get_param_names());
        }

        return detail::invoke_helper<Ret(Args...)>::invoke(
            *static_cast<Derived*>(this), args, Derived::get_param_names());
    }

    size_t get_arg_count() const override {
        return sizeof...(Args);
    }

    // 默认实现：返回通用参数名
    static constexpr std::array<const char*, sizeof...(Args)> get_param_names() {
        std::array<const char*, sizeof...(Args)> names{};
        for (size_t i = 0; i < sizeof...(Args); ++i) {
            names[i] = "lambda";
        }
        return names;
    }

    // 获取最小参数数量（考虑默认值）
    size_t get_min_arg_count() const override {
        if (defaults_.empty()) {
            return sizeof...(Args);
        }
        return first_default_index_;
    }

private:
    mc::variants defaults_;                             // 默认值数组
    size_t       first_default_index_{sizeof...(Args)}; // 第一个有默认值的参数索引
};

template <typename Ret = mc::variant, typename Lambda, typename... Args>
auto call(Lambda&& func, Args&&... args) -> std::enable_if_t<
    mc::is_variant_v<std::decay_t<Lambda>>, Ret> {
    if (!func.is_extension()) {
        MC_THROW(mc::runtime_exception, "对象不可调用: {type}",
                 ("type", func.get_type_name()));
    }

    auto ext        = func.as_extension();
    auto lambda_ext = mc::dynamic_pointer_cast<lambda_extension_base>(ext);

    if (!lambda_ext) {
        // 未来可在此处添加其他可调用类型的支持
        MC_THROW(mc::runtime_exception, "对象不可调用，实际类型: {type}",
                 ("type", ext->get_type_name()));
    }

    if constexpr (std::is_void_v<Ret>) {
        lambda_ext->call({std::forward<Args>(args)...});
    } else {
        return mcpy::cast<Ret>(lambda_ext->call({std::forward<Args>(args)...}));
    }
}

template <typename Ret = std::monostate, typename Lambda, typename... Args>
auto call(Lambda&& func, Args&&... args) -> std::enable_if_t<
    std::is_base_of_v<mcpy::lambda_extension_base, std::decay_t<Lambda>>,
    std::conditional_t<
        std::is_same_v<Ret, std::monostate>,
        typename std::decay_t<Lambda>::return_type, Ret>> {
    using return_type = std::conditional_t<
        std::is_same_v<Ret, std::monostate>,
        typename std::decay_t<Lambda>::return_type, Ret>;

    if constexpr (sizeof...(Args) == std::decay_t<Lambda>::arg_count) {
        if constexpr (std::is_void_v<return_type>) {
            func(std::forward<Args>(args)...);
        } else {
            return mcpy::cast<return_type>(func(std::forward<Args>(args)...));
        }
    } else {
        if constexpr (std::is_void_v<return_type>) {
            func.call_impl({std::forward<Args>(args)...});
        } else {
            return mcpy::cast<return_type>(func.call_impl({std::forward<Args>(args)...}));
        }
    }
}

} // namespace mcpy

#endif // MC_BUILTINS_LAMBDA_H
